#include <QtGui>
#include <QtCore>
#include "proxy_model.h"

//embedded class that stores with info on index mapping
//between the source model and it's proxy
class ProxyModel::Private {
public:
  ProxyModel * instance;
  QHash<int, int> sourceRootsCounts;
  bool aboutToRemoveRoots;

  QHash<int, int> subrow_counts;
  QHash<QPair<int, int>, int> subsubrow_counts;

  //count the number of rows corresponding to a row in the source model
  int source_root2count(int row) {
      if (!instance->sourceModel())
          return 0;

      if (!sourceRootsCounts.contains(row)) {
          sourceRootsCounts[row] = instance->sourceModel()->rowCount(
              instance->sourceModel()->index(row,0)
          );
      }

      return sourceRootsCounts[row];
  }


  int subrow_count(int row) {
    if (!instance->sourceModel())
      return 0;

    if (!subrow_counts.contains(row)) {
      subrow_counts[row] = instance->sourceModel()->rowCount(
        instance->sourceModel()->index(row,0)
      );
    }

    return subrow_counts[row];
  }

  int subsubrow_count(int row, int subrow) {
    if (!instance->sourceModel())
      return 0;

    QPair<int, int> pair = qMakePair(row, subrow);

    if (!subsubrow_counts.contains(pair)) {
      subsubrow_counts[pair] = instance->sourceModel()->rowCount(
        instance->sourceModel()->index(subrow, 0, instance->sourceModel()->index(row,0))
      );
    }

    return subsubrow_counts[pair];
  }

  int row_from_source(int row, int subrow, int subsubrow) {

    int destination_row = subsubrow;

    for(int r = 0; r < row - 1; ++r) {
      for(int s = 0; s < subrow_count(r); ++s) {
        destination_row += subsubrow_count(r,s);
      }
    }

    for(int s = 0; s < subrow; ++s) {
      destination_row += subsubrow_count(row-1,s);
    }

    return destination_row;
  }

  int row_fromSource(int root, int row) {
    for (int r = 0; r < root; r++)
      row += source_root2count(r);
    return row;
  }

  QVector<int> row_to_source(int flat_row) {
    QVector<int> ret;

    for (int row = 0; row < instance->sourceModel()->rowCount(); ++row) {
    //qDebug() << instance->sourceModel()->rowCount();

    for (int subrow = 0; subrow < subrow_count(row); ++subrow) {
      int subsubrows_in_subrow = subsubrow_count(row, subrow);
      //qDebug() << flat_row << subsubrows_in_subrow;
      if(flat_row >= subsubrows_in_subrow)
        flat_row -= subsubrows_in_subrow;
        else {
          ret << row << subrow << flat_row;
          //qDebug() << ret;
          return ret;
        }
      }
    }
    return ret;
  }

  QPair<int,int> row_toSource(int row) {
    int root = 0;
    for (int r =0; r < instance->sourceModel()->rowCount(); r++) {
      root = r;
      int rows_in_root = source_root2count(r);
      if (row >= rows_in_root)
        row -= rows_in_root;
        else break;
      }
      return qMakePair(root, row);
  }
};

ProxyModel::ProxyModel(QObject * parent)
:QAbstractProxyModel(parent)
{
    d = new Private;
    d->instance = this;
    d->aboutToRemoveRoots = false;
}

ProxyModel::~ProxyModel()
{
    delete d;
}

void ProxyModel::setSourceModel(QAbstractItemModel * m)
{
    if (sourceModel()) {
        disconnect(sourceModel(), SIGNAL(rowsAboutToBeInserted(QModelIndex, int, int)),
                   this, SLOT(source_rowsAboutToBeInserted(QModelIndex, int, int)));
        disconnect(sourceModel(), SIGNAL(rowsInserted(QModelIndex, int, int)),
                   this, SLOT(source_rowsInserted(QModelIndex, int, int)));

        disconnect(sourceModel(), SIGNAL(rowsAboutToBeRemoved(QModelIndex, int, int)),
                   this, SLOT(source_rowsAboutToBeRemoved(QModelIndex, int, int)));
        disconnect(sourceModel(), SIGNAL(rowsRemoved(QModelIndex, int, int)),
                   this, SLOT(source_rowsRemoved(QModelIndex, int, int)));

        disconnect(sourceModel(), SIGNAL(dataChanged(QModelIndex,QModelIndex)),
                   this, SLOT(source_dataChanged(QModelIndex,QModelIndex)));
        disconnect(sourceModel(), SIGNAL(modelReset()),
                   this, SLOT(source_modelReset()));

        disconnect(sourceModel(), SIGNAL(layoutAboutToBeChanged()), this, SIGNAL(layoutAboutToBeChanged()));
        disconnect(sourceModel(), SIGNAL(layoutChanged()), this, SIGNAL(layoutChanged()));
    }

    QAbstractProxyModel::setSourceModel(m);

    if (sourceModel()) {
        connect(sourceModel(), SIGNAL(rowsAboutToBeInserted(QModelIndex, int, int)),
                this, SLOT(source_rowsAboutToBeInserted(QModelIndex, int, int)));
        connect(sourceModel(), SIGNAL(rowsInserted(QModelIndex, int, int)),
                this, SLOT(source_rowsInserted(QModelIndex, int, int)));

        connect(sourceModel(), SIGNAL(rowsAboutToBeRemoved(QModelIndex, int, int)),
                this, SLOT(source_rowsAboutToBeRemoved(QModelIndex, int, int)));
        connect(sourceModel(), SIGNAL(rowsRemoved(QModelIndex, int, int)),
                this, SLOT(source_rowsRemoved(QModelIndex, int, int)));

        connect(sourceModel(), SIGNAL(dataChanged(QModelIndex,QModelIndex)),
                this, SLOT(source_dataChanged(QModelIndex,QModelIndex)));
        connect(sourceModel(), SIGNAL(modelReset()),
                this, SLOT(source_modelReset()));

        connect(sourceModel(), SIGNAL(layoutAboutToBeChanged()), this, SIGNAL(layoutAboutToBeChanged()));
        connect(sourceModel(), SIGNAL(layoutChanged()), this, SIGNAL(layoutChanged()));
    }

    //QAbstractItemModel::reset();
     beginResetModel();
     endResetModel();
}

QModelIndex ProxyModel::mapFromSource(const QModelIndex & source) const
{
    if (!sourceModel())
        return QModelIndex();

    if (!source.parent().isValid())
        return QModelIndex();
    //qDebug() << source.parent() << source.parent().parent();
    if (source.parent().parent().isValid()) {
        qDebug() << "row = " << d->row_from_source(source.parent().parent().row(), source.parent().row(), source.row());
    }

    //qDebug() << "row = " << d->row_fromSource(source.parent().row(), source.row()) << "column = " << source.column();
    return index(d->row_fromSource(source.parent().row(), source.row()), source.column());
}

//used for the initial setup
QModelIndex ProxyModel::mapToSource(const QModelIndex & proxy) const
{
    if (!sourceModel())
        return QModelIndex();

    //QPair<int, int> pos = d->row_toSource(proxy.row());
    //int root_row = pos.first;
    //int row = pos.second;

    QVector<int> position = d->row_to_source(proxy.row());
    //qDebug() << "row = " << proxy.row() << "corresponds to " << "position = " << position;
    int source_row = position[0];
    int source_subrow = position[1];
    int source_subsubrow = position[2];

    QModelIndex source_row_index = sourceModel()->index(source_row, 0);
    QModelIndex source_subrow_index = sourceModel()->index(source_subrow, 0, source_row_index);
    QModelIndex source_subsubrow_index = sourceModel()->index(source_subsubrow, proxy.column(), source_subrow_index);

    return source_subsubrow_index;

    //qDebug() << root_row;
    //QModelIndex p = sourceModel()->index(root_row, 0); //proxy.column()
    //qDebug() << p;//row << " " << proxy.column() << sourceModel()->index(row, proxy.column(), p);
    //return sourceModel()->index(row, proxy.column(), p);
}

QModelIndex	ProxyModel::parent(const QModelIndex &) const
{
    return QModelIndex();
}

//used for the initial setup
QModelIndex ProxyModel::index(int row, int column, const QModelIndex &) const
{
    return createIndex(row, column);
}

//used for the initial setup
int	ProxyModel::rowCount(const QModelIndex & p) const
{
    if (p.isValid()) return 0;
    if (!sourceModel()) return 0;

    int count = 0;
    for (int root_row =0; root_row< sourceModel()->rowCount(); root_row++) {
        for (int subrow = 0; subrow < d->subrow_count(root_row); ++subrow) {
            count += d->subsubrow_count(root_row, subrow); //d->source_root2count(root_row);
        }
    }

    //qDebug() << count;
    return count;
}

//used for the initial setup
int	ProxyModel::columnCount(const QModelIndex & p) const
{
    if (p.isValid()) return 0;
    if (!sourceModel()) return 0;
    return sourceModel()->columnCount();
}

void ProxyModel::source_rowsAboutToBeInserted(QModelIndex p, int from, int to)
{
    if (!p.parent().isValid())
        return;

    //int f = d->row_fromSource(p.row(), from);
    int f = d->row_from_source(p.parent().row(), p.row(), from);
    int t = f + (from-to);
    beginInsertRows(QModelIndex(), f, t);
}

void ProxyModel::source_rowsInserted(QModelIndex p, int, int)
{
    d->sourceRootsCounts.clear();
    if (!p.parent().isValid())
        return;

    endInsertRows();
}

void ProxyModel::source_rowsAboutToBeRemoved(QModelIndex p, int from, int to)
{
    if (!p.isValid()) {
        //remove root items
        int f = d->row_fromSource(from,0);
        int t = d->row_fromSource(to,0)+ d->source_root2count(to);

        if (f != t) {
            beginRemoveRows(QModelIndex(), f, t-1);
            d->aboutToRemoveRoots = true;
        }

        return;
    }

    int f = d->row_fromSource(p.row(), from);
    int t = f + (from-to);
    beginRemoveRows(QModelIndex(), f, t);
}

void ProxyModel::source_rowsRemoved(QModelIndex p, int, int)
{
    d->sourceRootsCounts.clear();

    if (!p.isValid()) {
        //remove root items
        if (d->aboutToRemoveRoots) {
            d->aboutToRemoveRoots = false;
            endRemoveRows();
        }
        return;
    }

    endRemoveRows();
}

void ProxyModel::source_dataChanged(QModelIndex tl, QModelIndex br)
{

    QModelIndex p_tl = mapFromSource(tl);
    QModelIndex p_br = mapFromSource(br);
    emit dataChanged(p_tl, p_br);
}

void ProxyModel::source_modelReset()
{
    beginResetModel();
    d->sourceRootsCounts.clear();
    endResetModel();
    //reset();
}

